import ecdsa
import numpy as np
import ssl
import asyncio
import websockets
import time
import random
from hashlib import sha256
from request import Request
from tvm_cifar10 import TvmDeployment, VGG16
from utils import cifar_test, cifar_train, get_dataset, cifar_test_loader, cifar_train_loader, attack_dataset
from multiprocessing import Process

admin_tls_dir = "/mnt/e/fabricdev/fabric-samples/test-network/organizations/peerOrganizations/org1.example.com/users/Admin@org1.example.com/tls"
ssl_context = ssl.create_default_context(
    ssl.Purpose.SERVER_AUTH, cafile=f"{admin_tls_dir}/ca.crt")
ssl_context.load_cert_chain(
    certfile=f"{admin_tls_dir}/client.crt", keyfile=f"{admin_tls_dir}/client.key")
ssl_context.verify_mode = ssl.CERT_REQUIRED
ssl_context.check_hostname = False


uri = "wss://localhost:8080/echo"
tvm_deployment = TvmDeployment()
nonces = [1]*10

# Normal Dataset
datas0, targets0 = attack_dataset(1, 1000, cifar_train_loader)
# Attack Dataset
datas, targets = attack_dataset(0.5, 1000, cifar_train_loader)


def priv_key(user_no):
    key_dir = f"/mnt/e/fabricdev/fabric-samples/test-network/organizations/peerOrganizations/org1.example.com/users/User{user_no}@org1.example.com/msp/keystore/priv_sk"
    return open(key_dir).read()


'''
websocket:  Websocket Connection.
deployment: ML deployment.
data:       Input tensor.
target:     True output.
nonce:      Tx serial number.
'''


async def send(websocket, data, target, nonce):

    # Build Dapp Request.
    req = Request()
    signing_key = ecdsa.SigningKey.from_pem(priv_key(target+1))
    req.sign(signing_key, sha256, data, nonce)

    # Upload Dapp Request.
    await websocket.send(req.with_sign)
    greeting = await websocket.recv()
    time.sleep(0.02)
    return nonce, greeting, target


async def register(websocket):

    hash = sha256()
    hash.update(bytes(tvm_deployment.lib.get_graph_json(), encoding="utf-8"))
    resp = await send(websocket, hash.digest(), 15, 0)
    print(resp, "Model Register End...")
    time.sleep(0.3)

    for i in range(10):
        idx = np.int64(i)
        hash = sha256()
        hash.update(idx.tobytes())
        resp = await send(websocket, hash.digest(), i, 0)
        print(resp, "User Register End...")
        time.sleep(0.3)
    pass


async def connect():
    async with websockets.connect(uri, ssl=ssl_context, ping_interval=None) as websocket:

        # Reigister Request
        await register(websocket)

        # Normal Request
        normal_counter = 0
        for i in range(len(datas0)):
            top1_tvm = tvm_deployment.fwd(datas0[i])
            hash = sha256()
            hash.update(top1_tvm.tobytes())
            resp = await send(websocket, hash.digest(), targets0[i],  nonces[targets0[i]])
            nonces[targets0[i]] += 1
            normal_counter += 1
            print(resp, top1_tvm)
            # if nonces[0] == 500:
            #     break

        # Attack Request
        for i in range(len(datas)):
            top1_tvm = tvm_deployment.fwd(datas[i])
            hash = sha256()
            hash.update(top1_tvm.tobytes())
            resp = await send(websocket, hash.digest(), targets[i], nonces[targets[i]])
            nonces[targets[i]] += 1
            print(resp, top1_tvm)

# rollup client
if __name__ == "__main__":

    asyncio.get_event_loop().run_until_complete(connect())
    # print(targets[:1000])
